from cmath import sqrt
from email.mime import image
import math
from types import FunctionType
from typing import Any, Dict, List, Tuple
import numpy
import cv2
from PIL import ImageFont, ImageDraw, Image
from PIL.ImageFont import FreeTypeFont
import torch
from word_vec_encoders.word_vec_encoder_base import WordVecEncoderBase


class GlyphsimpleWordVecEncoder(WordVecEncoderBase):
    """
    该编码器将字符渲染为小型图像并施加高斯模糊，结果直接用于特征
    """

    def __init__(self, font_file_name:str = 'SourceHanSansSC-Regular.otf', glyph_size:int = 12) -> None:
        super().__init__('glyphsimple', embedding_dim = glyph_size * glyph_size)
        self.font_file_name = font_file_name
        self.font_size = glyph_size - 2
        self.font = ImageFont.truetype(font_file_name, self.font_size)
        self.image_size = glyph_size
        self.embedding_dim = glyph_size * glyph_size

    def forward(self, input_tensor:torch.Tensor, batch_size:int):
        return input_tensor

    def collate_batch_tensor(self, batch_sentences:List[str]):
        input_tensor = self.batch_str_to_input_tensor_bert_style_cached(batch_sentences)
        batch_size, sentence_length, _ = input_tensor.size()        
        # 打破句子界限，组合为小序列的批次
        seq_batch_size = batch_size * sentence_length
        input_tensor = input_tensor.reshape((seq_batch_size, self.image_size, self.image_size, 1))
        return input_tensor

    def sentence_to_tensor(self, sentence:str) -> torch.Tensor:
        sentence_tensor = []
        for char in sentence:
            sentence_tensor.append(self.encode_char_cached(char, self.image_size, self.font))
        return torch.stack(sentence_tensor)

    def empty_chars_tensor(self, num_chars:int) -> torch.Tensor:
        return torch.zeros((num_chars, self.image_size * self.image_size))

    _glyph_cache:Dict[str, torch.Tensor] = dict()
    def encode_char_cached(self, char:str, image_size:int, font:FreeTypeFont):
        if char not in GlyphsimpleWordVecEncoder._glyph_cache:
            image = Image.fromarray(numpy.zeros((image_size, image_size), numpy.uint8))
            draw = ImageDraw.Draw(image)
            # 让字位于中间
            draw.text((0, 0),  char, fill ="white", font = font, align ="left")
            # 高斯模糊
            image = cv2.GaussianBlur(numpy.array(image), (3, 3), 0)
            image = torch.Tensor(image) / 255.0
            linear_image = image.reshape((self.embedding_dim))
            GlyphsimpleWordVecEncoder._glyph_cache[char] = linear_image
        glyph = GlyphsimpleWordVecEncoder._glyph_cache[char]
        return glyph